"""
测试数据收集器模块
"""
import os
import sys
import unittest
import numpy as np

# 添加项目根目录到路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..')))

from sichuanmajiang.data.data_collector import DataCollector, ExpertDataCollector, ExperienceCollector
from sichuanmajiang.tests.test_utils import create_test_game_state, create_test_directory, remove_test_directory


class TestDataCollector(unittest.TestCase):
    """
    测试DataCollector类
    """
    def setUp(self):
        """
        测试前的设置
        """
        self.test_dir = os.path.join(os.path.dirname(__file__), 'test_data')
        create_test_directory(self.test_dir)
        
        # 创建数据收集器
        self.collector = DataCollector(output_dir=self.test_dir)
    
    def tearDown(self):
        """
        测试后的清理
        """
        # 移除测试目录
        remove_test_directory(self.test_dir)
    
    def test_collect_state(self):
        """
        测试收集状态
        """
        # 创建测试游戏状态
        game_state = create_test_game_state()
        
        # 收集状态
        self.collector.collect_state(game_state)
        
        # 检查数据是否被收集
        self.assertEqual(len(self.collector.states), 1)
        self.assertEqual(self.collector.states[0]['current_player'], game_state['current_player'])
    
    def test_collect_action(self):
        """
        测试收集动作
        """
        # 收集动作
        action = 2  # 假设2表示'chi'
        self.collector.collect_action(action)
        
        # 检查数据是否被收集
        self.assertEqual(len(self.collector.actions), 1)
        self.assertEqual(self.collector.actions[0], action)
    
    def test_collect_reward(self):
        """
        测试收集奖励
        """
        # 收集奖励
        reward = 10.0
        self.collector.collect_reward(reward)
        
        # 检查数据是否被收集
        self.assertEqual(len(self.collector.rewards), 1)
        self.assertEqual(self.collector.rewards[0], reward)
    
    def test_save_to_numpy(self):
        """
        测试保存为NumPy格式
        """
        # 收集一些数据
        for i in range(5):
            game_state = create_test_game_state()
            game_state['game_round'] = i
            self.collector.collect_state(game_state)
            self.collector.collect_action(i % 6)
            self.collector.collect_reward(float(i))
        
        # 保存数据
        output_path = self.collector.save_to_numpy('test_data')
        
        # 检查文件是否存在
        self.assertTrue(os.path.exists(output_path))
        
        # 加载数据并检查
        loaded_data = np.load(output_path, allow_pickle=True)
        self.assertEqual(len(loaded_data['states']), 5)
        self.assertEqual(len(loaded_data['actions']), 5)
        self.assertEqual(len(loaded_data['rewards']), 5)
    
    def test_save_to_json(self):
        """
        测试保存为JSON格式
        """
        # 收集一些数据
        for i in range(3):
            game_state = create_test_game_state()
            game_state['game_round'] = i
            self.collector.collect_state(game_state)
            self.collector.collect_action(i % 6)
            self.collector.collect_reward(float(i))
        
        # 保存数据
        output_path = self.collector.save_to_json('test_data')
        
        # 检查文件是否存在
        self.assertTrue(os.path.exists(output_path))

    def test_reset(self):
        """
        测试重置收集器
        """
        # 收集一些数据
        for i in range(3):
            self.collector.collect_state(create_test_game_state())
            self.collector.collect_action(i % 6)
            self.collector.collect_reward(float(i))
        
        # 检查数据长度
        self.assertEqual(len(self.collector.states), 3)
        
        # 重置
        self.collector.reset()
        
        # 检查是否已重置
        self.assertEqual(len(self.collector.states), 0)
        self.assertEqual(len(self.collector.actions), 0)
        self.assertEqual(len(self.collector.rewards), 0)


class TestExpertDataCollector(unittest.TestCase):
    """
    测试ExpertDataCollector类
    """
    def setUp(self):
        """
        测试前的设置
        """
        self.test_dir = os.path.join(os.path.dirname(__file__), 'test_data')
        create_test_directory(self.test_dir)
        
        # 创建专家数据收集器
        self.collector = ExpertDataCollector(output_dir=self.test_dir)
    
    def tearDown(self):
        """
        测试后的清理
        """
        # 移除测试目录
        remove_test_directory(self.test_dir)
    
    def test_collect_expert_move(self):
        """
        测试收集专家动作
        """
        # 收集专家动作
        game_state = create_test_game_state()
        action = 2  # 'chi'
        confidence = 0.9
        
        self.collector.collect_expert_move(game_state, action, confidence)
        
        # 检查数据是否被收集
        self.assertEqual(len(self.collector.expert_data), 1)
        self.assertEqual(self.collector.expert_data[0]['action'], action)
        self.assertEqual(self.collector.expert_data[0]['confidence'], confidence)
    
    def test_collect_game_record(self):
        """
        测试收集游戏记录
        """
        # 创建游戏记录
        game_record = {
            'players': ['player1', 'player2', 'player3', 'player4'],
            'winner': 'player1',
            'scores': [2000, 500, 500, 0],
            'moves': 120,
            'duration': 300
        }
        
        # 收集游戏记录
        self.collector.collect_game_record(game_record)
        
        # 检查数据是否被收集
        self.assertEqual(len(self.collector.game_records), 1)
        self.assertEqual(self.collector.game_records[0]['winner'], 'player1')
    
    def test_save_expert_data(self):
        """
        测试保存专家数据
        """
        # 收集一些专家数据
        for i in range(2):
            game_state = create_test_game_state()
            game_state['game_round'] = i
            self.collector.collect_expert_move(game_state, i % 6, 0.8 + i * 0.1)
        
        # 保存数据
        output_path = self.collector.save_expert_data('expert_data')
        
        # 检查文件是否存在
        self.assertTrue(os.path.exists(output_path))


class TestExperienceCollector(unittest.TestCase):
    """
    测试ExperienceCollector类
    """
    def setUp(self):
        """
        测试前的设置
        """
        # 创建经验收集器
        self.collector = ExperienceCollector(buffer_size=100)
    
    def test_add_experience(self):
        """
        测试添加经验
        """
        # 添加经验
        state = np.array([1.0, 2.0, 3.0])
        action = 1
        reward = 0.5
        next_state = np.array([1.1, 2.1, 3.1])
        done = False
        
        self.collector.add_experience(state, action, reward, next_state, done)
        
        # 检查经验是否被添加
        self.assertEqual(len(self.collector.buffer), 1)
    
    def test_sample_batch(self):
        """
        测试采样批次
        """
        # 添加一些经验
        for i in range(20):
            state = np.array([i, i+1, i+2], dtype=np.float32)
            action = i % 6
            reward = float(i) * 0.1
            next_state = np.array([i+0.1, i+1.1, i+2.1], dtype=np.float32)
            done = i == 19
            
            self.collector.add_experience(state, action, reward, next_state, done)
        
        # 采样批次
        batch_size = 5
        batch = self.collector.sample_batch(batch_size)
        
        # 检查批次大小
        self.assertEqual(len(batch['states']), batch_size)
        self.assertEqual(len(batch['actions']), batch_size)
        self.assertEqual(len(batch['rewards']), batch_size)
        self.assertEqual(len(batch['next_states']), batch_size)
        self.assertEqual(len(batch['dones']), batch_size)
    
    def test_get_buffer_size(self):
        """
        测试获取缓冲区大小
        """
        # 初始大小应为0
        self.assertEqual(self.collector.get_buffer_size(), 0)
        
        # 添加一些经验
        for i in range(10):
            state = np.array([i, i+1, i+2], dtype=np.float32)
            action = i % 6
            reward = float(i) * 0.1
            next_state = np.array([i+0.1, i+1.1, i+2.1], dtype=np.float32)
            done = False
            
            self.collector.add_experience(state, action, reward, next_state, done)
        
        # 检查大小
        self.assertEqual(self.collector.get_buffer_size(), 10)
    
    def test_reset_buffer(self):
        """
        测试重置缓冲区
        """
        # 添加一些经验
        for i in range(5):
            state = np.array([i, i+1, i+2], dtype=np.float32)
            action = i % 6
            reward = float(i) * 0.1
            next_state = np.array([i+0.1, i+1.1, i+2.1], dtype=np.float32)
            done = False
            
            self.collector.add_experience(state, action, reward, next_state, done)
        
        # 重置缓冲区
        self.collector.reset_buffer()
        
        # 检查大小
        self.assertEqual(self.collector.get_buffer_size(), 0)
    
    def test_buffer_capacity(self):
        """
        测试缓冲区容量限制
        """
        # 创建一个小容量的收集器
        small_collector = ExperienceCollector(buffer_size=5)
        
        # 添加超过容量的经验
        for i in range(10):
            state = np.array([i, i+1, i+2], dtype=np.float32)
            action = i % 6
            reward = float(i) * 0.1
            next_state = np.array([i+0.1, i+1.1, i+2.1], dtype=np.float32)
            done = False
            
            small_collector.add_experience(state, action, reward, next_state, done)
        
        # 检查大小（应该等于容量）
        self.assertEqual(small_collector.get_buffer_size(), 5)


if __name__ == '__main__':
    unittest.main()
